#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2019/5/16 15:22
# @Author  : CHEN Wang
# @Site    :
# @File    : hammer.py
# @Software: PyCharm

import multiprocessing
import random
import pymysql
import pandas as pd
import numpy as np
import math
import prettytable
from quant_researcher.quant.project_tool.logger.my_logger import LOG, PRINT_KEY_INFORMATION
from quant_researcher.quant.project_tool.pretty_table import df_to_table
from quant_researcher.quant.project_tool.parallelism import MultiThread
from quant_researcher.quant.project_tool import assert_tool, celebrity
from quant_researcher.quant.project_tool.db_operator import db_conn, conn_dispatcher


def split_a_list_into_batches(a_list, how_many):
    """
    把一个 list 分成几份，尽量均分

    :raise ValueError: 如果传入的份数 = 0
    :param a_list: 数据
    :param int how_many: 几份
    :return: list，每一个 element 都是源数据的子集
    """
    # 2019年12月22日14:22:21，如果这个数组为空的话怎么办？为了保持一致性
    # 因为这个函数的结果一般是用来进行循环的，所以返回一个空数组看起来是比较合理的
    if not a_list:
        LOG.warning(f'传入的数组判断为False：{a_list}，为了保持结果的一致性，返回[]')
        return []

    if how_many == 0:
        raise ValueError('需要分成的进程数不能是 0')
    total_count = len(a_list)
    LOG.info(f'总元素个数：{total_count}，希望进程数：{how_many}')
    actual_pool_size = min(total_count, how_many)
    step_max = math.ceil(total_count / actual_pool_size)
    step_min = math.floor(total_count / actual_pool_size)
    batch_list_size = [step_min] * actual_pool_size
    gap = total_count - (step_min * actual_pool_size)
    batch_list_size[:gap] = [step_max] * gap
    LOG.info(f'{actual_pool_size}个进程, 每个进程分到的元素个数：{batch_list_size}')

    start_idx = 0
    batch_list = []
    for idx in batch_list_size:
        end_idx = start_idx + idx
        batch_list.append(a_list[start_idx:end_idx])
        start_idx = end_idx
    assert sum([len(x) for x in batch_list]) == total_count

    return batch_list


def _deal_a_batch_with_a_loop(a_batch, func, **kwargs):
    """
    循环将一个 list 中的元素传入指定函数运行，并将每一次运行的结果放入 list 返回; 这里是将所有任务分成多个batch后，对一个batch内的多任务进行循环处理
    传入每个进程时需要提前创建好数据库连接，每个进程共用相同的数据库连接

    :param list a_batch: 数据
    :param func: 函数
    :param kwargs:
        - set_basic_data_conn，bool，建立一个到基础数据库的 conn，默认：False
        - set_derivative_data_conn，bool，建立一个到衍生指标库的 conn，默认：False
        - set_stock_conn，bool，建立一个到A股数据库的 conn，默认：False
        - set_third_party_conn，bool，建立一个到维赛特数据库的 conn，默认：False
        - set_fof_conn，bool，建立一个到基金组合数据库的 conn，默认：False
        - by_sqlalchemy，bool，数据库连接是否用create_engine方式创建
        - ignore_err，bool，忽略某一次运行的发生的错误，默认：False
    :return: list
    """
    set_basic_data_conn = kwargs.pop('set_basic_data_conn', False)
    set_derivative_data_conn = kwargs.pop('set_derivative_data_conn', False)
    set_derivative_data_write_conn = kwargs.pop('set_derivative_data_write_conn', False)
    set_stock_conn = kwargs.pop('set_stock_conn', False)
    set_third_party_conn = kwargs.pop('set_third_party_conn', False)
    set_fof_conn = kwargs.pop('set_fof_conn', False)
    by_sqlalchemy = kwargs.pop('by_sqlalchemy', False)
    ignore_err = kwargs.pop('ignore_err', False)

    conn_dict = {}
    if set_basic_data_conn:
        conn_dict.update({'conn': db_conn.get_basic_data_conn()})
    if set_derivative_data_conn:
        conn_dict.update({'derivative_data_conn': db_conn.get_derivative_data_conn()})
    if set_derivative_data_write_conn:
        conn_dict.update({'derivative_data_write_conn': db_conn.get_derivative_data_write_conn()})
    if set_stock_conn:
        conn_dict.update({'stock_conn': db_conn.get_stock_conn()})
    if set_third_party_conn:
        conn_dict.update({'third_party_conn': db_conn.get_third_party_conn()})
    if set_fof_conn:
        conn_dict.update({'fof_conn': db_conn.get_fof_conn(by_sqlalchemy=by_sqlalchemy)})
    kwargs.update(conn_dict)

    def close_conn_dict(a_dict):
        """
        关闭一个连接字典

        :param a_dict: 连接字典
        :return:
        """
        for conn_name in a_dict:
            if type(a_dict[conn_name]) == pymysql.connections.Connection:
                LOG.info(f'现在关闭连接：{conn_name}')
                a_dict[conn_name].close()

    res = []
    for x in a_batch:
        try:
            res.append(func(x, **kwargs))
        except Exception as err:
            LOG.error(f'处理 {x} 时发生错误：{err}')
            if not ignore_err:
                close_conn_dict(conn_dict)
                raise err
    close_conn_dict(conn_dict)
    LOG.info(f'该进程任务处理结束')
    return res


def _deal_a_batch_with_a_loop_v2(a_batch, func, **kwargs):
    """
    循环将一个 list 中的元素传入指定函数运行，并将每一次运行的结果放入 list 返回； ; 这里是将所有任务分成多个batch后，对一个batch内的多任务进行循环处理

    与上述区别在于传入每个进程时，不需要提前创建好数据库连接，只需要制定数据库连接名
    这样的好处时，每个进程使用的数据库连接是分开的

    :param list a_batch: 数据
    :param func: 函数
    :param kwargs:
        - prepare_these_conn，list，需准备的 conn 的名字，默认：[]
        - ignore_err，bool，忽略某一次运行的发生的错误，默认：False
    :return: list
    """
    prepare_these_conn = kwargs.pop('prepare_these_conn', [])
    ignore_err = kwargs.pop('ignore_err', False)
    # LOG.info(f'需准备的连接：{prepare_these_conn}，是否忽略错误：{ignore_err}')

    conn_dict = None
    if prepare_these_conn:
        conn_dict = conn_dispatcher.dispatcher(db_names=prepare_these_conn)
        kwargs.update({'conn_dict': conn_dict})

    res = []
    for x in a_batch:
        try:
            res.append(func(x, **kwargs))
        except Exception as err:
            LOG.error(f'处理 {x} 时发生错误：{err}')
            if not ignore_err:
                conn_dispatcher.close_them(conn_dict)
                raise err
    conn_dispatcher.close_them(conn_dict)
    LOG.info(f'该进程任务处理结束')
    return res


def batch_process(a_list, func, **kwargs):
    """
    批量处理，把 list 中的每个对象都交给 func 处理并附上指定的关键字参数 info； 这里的list是已经进行分割好的batch

    :param a_list: 其中每一个对象都会交给func处理
    :param func: 处理的函数
    :param kwargs:
        - multiple_threads，bool，是否使用多线程，默认：True
        - info，dict，将会作为关键字参数传递给func，默认：{}
        - step，int，如果是多线程处理，每次最大的线程数，默认：30
        - ignore_err: bool，是否忽略错误，默认：False
        - pool_size: int，目标函数执行的最多进程数，默认：None
        - verbose，bool，是不是把函数执行的第一个参数也一并使用 tuple 返回
            方便在多进程和多线程拿到结果的时候对比，要不然是乱的，默认：False
    :return:
        list，每一个对象时每次func运行的返回值
    """
    ignore_err = kwargs.pop('ignore_err', False)
    multiple_threads = kwargs.pop('multiple_threads', True)
    verbose = kwargs.pop('verbose', False)

    assert_tool.value_is_instance(a_list, list)

    total_count = len(a_list)
    LOG.info(f'输入数组元素个数={total_count}，verbose={verbose}，'
             f'是否多线程={multiple_threads}，处理函数：{func.__name__}')

    # 2019年12月22日14:22:21，如果这个数组为空的话怎么办？为了保持一致性
    # 因为这个函数的结果一般是用来进行循环的，所以返回一个空数组看起来是比较合理的
    if not a_list:
        LOG.warning(f'传入的数组判断为False：{a_list}，为了保持结果的一致性，返回[]')
        return []

    total_res = []
    if multiple_threads:
        LOG.info('将采用多线程来执行任务')
        step = kwargs.pop('step', 30)
        LOG.info('多线程时的最大线程同时数量: %s', step)

        def for_a_batch(batch, info):
            """deal with a batch"""
            res = []
            thread_list = []
            for a_code in batch:
                # 输入总是传给func的第一个参数
                a_thread = MultiThread(func, args=(a_code,), info=info)
                a_thread.start()
                thread_list.append((a_code, a_thread))
            for item in thread_list:
                a_thread = item[1]
                a_thread.join()
                a_res = a_thread.get_result()
                if verbose:
                    res.append((item[0], a_res))
                else:
                    res.append(a_res)
            return res

        start_idx = 0
        while start_idx < total_count:
            end_idx = min(start_idx + step, total_count)
            this_batch = a_list[start_idx:end_idx]
            total_res.extend(for_a_batch(this_batch, kwargs.get('info', {})))
            start_idx = end_idx
    else:
        pool_size = kwargs.pop('pool_size', None)
        if pool_size is None:
            raise ValueError('传入的 pool size 是 None，或者根本没传')
        pool_size = min(pool_size, total_count)
        LOG.info(f'将采用多进程来执行任务！'
                 f'任务数量：{total_count}，进程池大小：{pool_size}')
        pool = multiprocessing.Pool(pool_size)
        for the_one in a_list:
            # 理论上，Python是不能pickle函数的
            # 只能将函数名和模块名进行pickle，所以就不能给函数加wrapper
            # apply_async(
            #     func, args=(), kwds={}, callback=None, error_callback=None
            # )
            kwds = kwargs.get('info', {})
            kwds.update({'ignore_err': ignore_err})
            total_res.append(pool.apply_async(func, args=(the_one,), kwds=kwds))
        LOG.warning('所有数据已加入进程池，即将关闭进程池')
        pool.close()
        LOG.warning('已关闭进程池，等待所有进程结束 ...')
        pool.join()
        LOG.warning(f'所有进程都已结束！已获得结果数量：{len(total_res)}')
        total_res = [r.get() for r in total_res]

    assert len(total_res) == len(a_list)
    return total_res


def split_a_list_into_batches_and_deal_each_with_a_process(the_list, **kwargs):
    """
    将一个 list 分批，采用多进程并行处理每一批，并把所有的运行结果放入一个 list 返回

    :raise ValueError: 如果 what_version 不是 v1 或 v2
    :param the_list: 数据
    :param kwargs:
        - how_many，int，分成的批数
        - info，dict，传给单体处理函数的关键字参数
        - ignore_err，bool，是否忽略错误（针对个体运行时的），默认：False
        - shuffle_the_list，bool，是否打乱数组中元素的顺序，默认：False
        - func_2_run，function，数组分批后，如何处理每一批的函数，默认：None
        - what_version，str，使用函数的第几个版本的功能，默认：v1
    :return: list，包含了每一个个体被处理后的结果
    """
    how_many = kwargs.pop('how_many')
    info = kwargs.pop('info')
    ignore_err = kwargs.pop('ignore_err', False)
    shuffle_the_list = kwargs.pop('shuffle_the_list', False)
    func_2_run = kwargs.pop('func_2_run', None)
    what_version = kwargs.pop('what_version', 'v1')

    # 2019年12月22日14:22:21，如果这个数组为空的话怎么办？为了保持一致性
    # 因为这个函数的结果一般是用来进行循环的，所以返回一个空数组看起来是比较合理的
    if not the_list:
        LOG.warning(f'传入的数组，False：{the_list}，为了保持结果的一致性，返回[]')
        return []

    if shuffle_the_list:
        random.shuffle(the_list)  # 在原list上进行操作，修改了原对象
    batches = split_a_list_into_batches(the_list, how_many)

    if func_2_run is None:
        if what_version == 'v1':
            func_2_run = _deal_a_batch_with_a_loop
        elif what_version == 'v2':
            func_2_run = _deal_a_batch_with_a_loop_v2
        else:
            raise ValueError('what_version 不是 v1 或 v2')

    batch_res = batch_process(
        batches, func=func_2_run, info=info,
        ignore_err=ignore_err, multiple_threads=False, pool_size=how_many
    )

    res = []
    for a_piece in batch_res:
        res.extend(a_piece)
    return res


def debug_dict(obj, name=None, a_f=False):
    """
    优雅地打印一个 dict 对象

    :raise RuntimeError: 如果传入的对象不是 dict
    :param dict obj: 数据
    :param str name: 名字，方便区分&搜索，默认：None
    :param bool a_f: 打印完是否结束程序，默认：False（结束程序）；True（继续运行）
    :return:
    """
    if not isinstance(obj, dict):
        raise RuntimeError
    x_table = prettytable.PrettyTable(['键', '值'])
    if name:
        x_table.add_row(['Name', name])
    for k, v in obj.items():
        x_table.add_row([k, v])
    is_all_number = True
    for x in obj.values():
        if not isinstance(x, int) and not isinstance(x, float):
            is_all_number = False
    if is_all_number:
        x_table.add_row(['V之和', sum(obj.values())])
    print(x_table)
    assert a_f


def debug_df(df, **kwargs):
    """
    优雅地打印一个 pd.DataFrame

    :param df: 数据
    :param kwargs:
        - a_f，bool，打印完是否结束程序，默认：False（结束程序）；True（继续运行）
        - show_all，bool，是不是将整个 df 全部打印，默认：False
        - deal_float，bool，是否需要处理 float 型数据，默认：True
        - show_index，bool，是否打印 df 的 index，默认：False
        - show_columns，bool，是否打印 df 有哪些列，默认：False
    :return:
    """
    digit = kwargs.pop('digit', 4)
    name = kwargs.pop('name', None)
    a_f = kwargs.pop('a_f', False)
    show_all = kwargs.pop('show_all', False)
    show_columns = kwargs.pop('show_columns', False)
    show_index = kwargs.pop('show_index', False)
    deal_float = kwargs.pop('deal_float', True)

    name = '*' if name is None else name
    my_df = df.copy() if show_index else df

    if my_df is None:
        print(f'这个DF< {name} >是None')
    else:
        how_many_markers = 15
        if_quit = '继续' if a_f else '退出'
        if_deal_float = f'🍀={digit}' if deal_float else '🍁'
        len_df = my_df.shape[0]
        msg = f'{name}--行数：{len_df}--{if_deal_float}-- 👉 {if_quit} '
        sepper = '-' * how_many_markers + '--S--' + msg + '--S--' + '-' * how_many_markers
        print(sepper)

        if show_index:
            idx_me = my_df.index.name
            print(f'此DF的索引名字：{idx_me}')
            my_df.insert(0, 'IDX' if idx_me is None else idx_me, my_df.index)

        if show_columns:
            cols = my_df.columns
            len_cols = len(my_df.columns)
            if len_cols <= 5:
                print(f'< {name} >的列名(%s)：%s, 值：%s' % (len_cols, cols.name, cols.values))
            elif 5 <= len_cols <= 10:
                print(f'< {name} >的列名(%s)：%s, 值：%s' % (len_cols, cols.name, list(cols.values)))
            else:
                print(f'< {name} >的列名(%s)：%s, 值:' % (len_cols, cols.name))
                print(my_df.columns.values)

        # print(f'debug_df::是否处理小数：{deal_float}，小数位数指定：{digit}')

        # 若没有特殊说明或特殊说明打印全部，则显示全部内容
        if show_all or len_df <= 31:
            print(df_to_table(my_df, deal_float, digit=digit))

        # 若没有特殊说明，只打印前后10行
        else:
            n_row = 10
            print(f'前%s行：' % n_row)
            print(df_to_table(my_df.head(n_row), deal_float, digit=digit))
            print(f'后%s行：' % n_row)
            print(df_to_table(my_df.tail(n_row), deal_float, digit=digit))
        sepper = '-' * how_many_markers + '--E--' + msg + '--E--' + '-' * how_many_markers
        print(sepper)  # 为什么不直接打印带 E 的一行呢，有的时候如果打印了很多，需要跳到结尾，直接搜索 name 便 OK
    assert a_f


def debug_matrix(item, name):
    """
    优雅地打印一个矩阵

    :param np.ndarray item: 数据
    :param str name: 名字
    :return:
    """
    name = 'unknown' if name is None else name

    if not isinstance(item, np.ndarray):
        print(f'非ndarray，{name}({item.__class__}): {item}\n')
        return

    max_col_or_row = 10  # 行或列超过这么多，再使用表格打印就看不清楚了，还不如直接使用它自己的呢
    on_each_side = int(max_col_or_row / 2)
    digit = 8  # 毕竟 numpy 自己的函数都只是显示 8 位
    column_marker, row_marker = 'C', 'R'
    shape = item.shape
    len_shape = len(shape)
    # print(f'shape of the item: {shape}, len of the shape: {len_shape}')
    if len_shape == 1:
        shape_str = shape
        single_row = []
        if shape[0] <= max_col_or_row:
            cols = [str(x + 1) for x in range(shape[0])]
            single_row.extend([celebrity.to_fixed_digit(x, digit=digit) for x in item])
        else:
            cols = []
            marker_added = False
            for col_idx in range(shape[0]):
                if on_each_side < col_idx + 1 <= shape[0] - on_each_side:
                    if not marker_added:
                        cols.append('...')
                        single_row.append('...')
                        marker_added = True
                else:
                    cols.append(str(col_idx + 1))
                    single_row.append(celebrity.to_fixed_digit(item[col_idx], digit=digit))
        x_table = prettytable.PrettyTable(cols)
        x_table.add_row(single_row)
    elif len_shape == 3:
        sepper = '-' * 22 + '这种矩阵目前不在我的考虑中' + '-' * 22
        print(sepper)
        print(item)
        print(sepper)
        return
    else:
        shape_str = '%s×%s' % (shape[0], shape[1])
        if shape[1] <= max_col_or_row:
            cols = [column_marker + str(x) if x != 0 else '' for x in range(shape[1] + 1)]
        else:
            cols = ['']
            marker_added = False
            for col_idx in range(shape[1]):
                if on_each_side < col_idx + 1 <= shape[1] - on_each_side:
                    if not marker_added:
                        cols.append('...')
                        marker_added = True
                else:
                    cols.append(column_marker + str(col_idx + 1))
        # print(shape_str)
        # print(cols)
        x_table = prettytable.PrettyTable(cols)

        marker_added = False
        for idx in range(shape[0]):
            single_row = []
            if on_each_side < idx + 1 < shape[0] - on_each_side + 1:
                if not marker_added:
                    marker_added = True
                    single_row.append(':')
                    single_row.extend([
                        ':' if x != on_each_side else '...' for x in range(min(on_each_side * 2 + 1, shape[1]))
                    ])
            else:
                single_row.append(f'{row_marker}{idx+1}')
                col_marker_added = False
                for col_idx in range(shape[1]):
                    if on_each_side < col_idx + 1 <= shape[1] - on_each_side:
                        if not col_marker_added:
                            single_row.append('...')
                            col_marker_added = True
                    else:
                        single_row.append(celebrity.to_fixed_digit(item[idx, col_idx], digit=digit))
            # print(single_row)
            if single_row:
                x_table.add_row(single_row)

    how_many_markers = 15
    msg = f'{name}--Shape={shape_str}'
    sepper = '-' * how_many_markers + '< S >' + msg + '< S >' + '-' * how_many_markers
    print(sepper)
    print(x_table if x_table else item)
    sepper = '=' * how_many_markers + '< E >' + msg + '< E >' + '=' * how_many_markers
    print(sepper)


def slim(*args, **kwargs):
    """
    打印出关键的中间结果
        直接或把他们保存下来之后，可以方便检查运行的步骤是不是符合预期。
        这样可以不用去数据库创建中间结果表，也可以达到debug的目的。
        那么一定要打印的很好看，方便阅读。
        为什么要放在这个模块：感觉这个函数的作用和LOG的作用差不多，而且都还需要从配置文件读取相关配置。
        记住：只打印帮助debug的关键信息，太多的输出只会带来负面作用
    :param args: 位置参数
    :param kwargs:
        - name，str，如果是其他类型对象的话，这个对象的名字，默认：None
        - digit，int，显示 pd.DataFrame 的时候，如果需要处理 float，保留小数位数，默认：4
        - a_f，bool，assert function的缩写，默认：False
        - show_all，bool，是不是显示 pd.DataFrame 的所有数据，默认：False
        - show_columns，bool，是不是显示 pd.DataFrame 的列信息，默认：False
        - show_index，bool，是不是显示 pd.DataFrame 的索引信息，默认：False
        - deal_float，bool，是不是处理 pd.DataFrame 中的float数据，方便显示，默认：True
        - print_anyway，bool，是不是不管怎样都要进行打印，默认：False
            有时配置文件中配置了不显示，但是又有特殊情况，比如说在 try catch 中，2020年3月24日17:54:37
    :return:
    """
    # print(locals())
    print_anyway = kwargs.pop('print_anyway', False)

    if PRINT_KEY_INFORMATION or print_anyway:
        name = kwargs.pop('name', None)
        digit = kwargs.pop('digit', 4)
        deal_float = kwargs.pop('deal_float', True)
        show_index = kwargs.pop('show_index', False)
        show_all = kwargs.pop('show_all', False)
        show_columns = kwargs.pop('show_columns', False)
        a_f = kwargs.pop('a_f', False)

        if len(args) == 1 and isinstance(args[0], pd.DataFrame):
            # 需要打印的对象为1个且为pd.DataFrame
            debug_df(args[0], a_f=a_f, show_all=show_all,
                     show_columns=show_columns, show_index=show_index,
                     deal_float=deal_float, name=name, digit=digit)

        elif len(args) == 1 and isinstance(args[0], np.ndarray):
            debug_matrix(args[0], name=name)

        elif len(args) == 1 and isinstance(args[0], dict):
            debug_dict(args[0], name=name, a_f=a_f)

        elif len(args) == 1 and isinstance(args[0], list):
            len_obj = len(args[0])
            if name:
                print(f'List< {name} -LEN：{len_obj}->的详情如下：↓')
            if len_obj > 100:
                how_many_to_print = 10
                for x in range(0, how_many_to_print):
                    print('No.', x+1, ', ', args[0][x], sep='')
                for x in range(len_obj - how_many_to_print, len_obj):
                    print('No.', x+1, ', ', args[0][x], sep='')
            else:
                for x in range(len_obj):
                    mull = args[0][x]
                    if isinstance(mull, pd.DataFrame) or isinstance(mull, dict) or isinstance(mull, np.ndarray):
                        slim(mull, name=f'No.{x + 1}, {name}', **{'a_f': 1})
                    else:
                        print('No.', x+1, ', ', mull, sep='')
        else:
            print(*args, **kwargs)


if __name__ == '__main__':
    pass
